#!/usr/bin/env python

import CSL as cs
import os
import sys
import math
import ConfigParser
from optparse import OptionParser
from string import Template
import random
import Queue
import threading



parser = OptionParser()
parser.add_option("-i", "--init", dest="initialization",
                  default='StringEvolutionT.ini',
                  help="set initialization FILE", metavar="FILE")
(opt,args) = parser.parse_args()

class Parameters:
    def __init__(self,initfile):
        self.config = ConfigParser.ConfigParser()
        self.config.read(initfile)
        try:
            self.NX  = self.config.getint('SIMULATION','LatticeSizeX')
            self.NY  = self.config.getint('SIMULATION','LatticeSizeY')
            self.NZ  = self.config.getint('SIMULATION','LatticeSizeZ')
            self.seed = self.config.getint('SIMULATION','RandomSeed')
            if self.seed == 0:
                rnd = random.SystemRandom()
                self.seed = rnd.randint(1,999999999)
            self.branches = self.config.getint('SIMULATION','ModelN')
            self.OM = self.config.getfloat('SIMULATION','OmegaMatter')
            self.OR = self.config.getfloat('SIMULATION','OmegaRadiation')
            self.h = self.config.getfloat('SIMULATION','LatticeSpacing')
            self.dt = self.config.getfloat('SIMULATION','TimeStep')
            self.a0 = self.config.getfloat('SIMULATION','Scale0')
            self.H0 = self.config.getfloat('SIMULATION','H0')
            self.maxHorizon = self.config.getfloat('SIMULATION','MaxHorizon')
            self.TasksX = self.config.getint('SIMULATION','TasksX')
            self.TasksY = self.config.getint('SIMULATION','TasksY')
            self.TasksZ = self.config.getint('SIMULATION','TasksZ')
            self.MeasureStep = self.config.getfloat('SIMULATION','MeasureStep')
            self.StartFromCheckpoint = self.config.getboolean('SIMULATION','StartFromCheckpoint')
            self.InitialMomentumSigma = self.config.getfloat('SIMULATION','InitialMomentumSigma')
            self.tpl = self.config.get('OUTPUT','DataTemplate')
            self.header = self.config.get('OUTPUT','DataHeader')
            self.outputdat = self.config.get('OUTPUT','DataFile')
            self.screenout = self.config.getboolean('OUTPUT','TalkOnScreen')
            self.checkpointfilename = self.config.get('OUTPUT','CheckpointFileName')
            self.checkpointinterval = self.config.getint('OUTPUT','CheckpointInterval')
            self.animationsave = self.config.getboolean('ANIMATION','AnimationSave')
            self.animationstep = self.config.getint('ANIMATION','AnimationStep')
            self.animationfile = self.config.get('ANIMATION','AnimationFile')
        except ConfigParser.NoSectionError,err:
            print "%s" % (err)            
            raise SystemExit
        except ConfigParser.NoOptionError,err:
            print "%s" % (err)            
            raise SystemExit

par = Parameters(opt.initialization)
if (par.NX % par.TasksX != 0) or (par.NY % par.TasksY != 0) or (par.NZ % par.TasksZ != 0):
    print "Lattice can't be segmented in the proposed way"
    raise SystemExit

data_template = Template(par.tpl)

cs.init_genrand64(par.seed)

latt1 = cs.latticeSP(par.NX,par.NY,par.NZ,par.branches)
latt2 = cs.latticeSP(par.NX,par.NY,par.NZ,par.branches)
latt1.init_random(par.InitialMomentumSigma)
frw       = cs.FriedmannEvolution(par.OR,par.OM,par.a0,par.H0)

if par.StartFromCheckpoint:
    print "START FROM CHECKPOINT"
    f = cs.InputFile(par.checkpointfilename)
    latt1.load(f.file())
    frw.load(f.file())
    f = None


evolution = cs.evolutionSP(par.h,par.dt,par.NX,par.NY,par.NZ,par.TasksX,par.TasksY,par.TasksZ)
measure   = cs.densitySP(par.NX,par.NY,par.NZ,par.TasksX,par.TasksY,par.TasksZ)

data_output = open(par.outputdat,'wb')
if par.screenout:
    sys.stdout.write(par.header+'\n')
data_output.write(par.header+'\n')
data_output.flush()
data_dict = {}

k = 0
p = 0

invH = 1.0/frw.get_H()
a = frw.get_scale_parameter()
t = frw.get_eta()
next_measure = t-1.0
sourcefirst = True

dtau = par.dt
old_horizon = 0.0
new_horizon = 0.0

while (invH/a)<par.maxHorizon:
    if (k % par.checkpointinterval)==0:
        print "Checkpointing: k = ",k
        f = cs.OutputFile(par.checkpointfilename)
        latt1.save(f.file())
        frw.save(f.file())
        f = None
        data_output.write('#-------------------CHECKPOINT----------------------\n')
        data_output.flush()

    if par.animationsave:
        if (k % par.animationstep)==0: 
            latt1.save_strings_vtk(par.animationfile % (p) )
            p = p + 1
    if sourcefirst:
        evolution(latt1,latt2,frw)
        sourcefirst = False
    else:
        evolution(latt2,latt1,frw)
        sourcefirst = True

    old_horizon  = new_horizon
    new_horizon = invH/a
    
    delta = new_horizon-old_horizon
    if delta>0.001:
        dtau = dtau/1.1
    else:
        dtau = dtau*1.1
    #print delta,dtau    
    evolution.set_dt(dtau)
    t = frw.get_eta()

    if t>=next_measure:
        measure(latt1)
        invH = 1.0/frw.get_H()
        a = frw.get_scale_parameter()
        data_dict["ctime"]  = "%g" % (frw.get_eta())
        data_dict["density"] = "%g" % (measure.total_density())
        data_dict["a"]       = "%g" % (a)
        data_dict["invH"]    = "%g" % (invH)
        data_dict["dtau"]    = "%g" % (dtau)
        msg = data_template.substitute(data_dict)
        lst = eval(msg)
        dtxt = ''
        for v in lst:
            dtxt = dtxt + ("%.5g, " % (v))
        dtxt = dtxt + '\n'
        if par.screenout:
            sys.stdout.write(dtxt)
        data_output.write(dtxt)
        data_output.flush()
        next_measure = t + par.MeasureStep

    k = k + 1
    if( (k % 10)==0):
        pass
        #cmd = """curl -d "data=%s&k=%d" http://localhost:8080""" % (dtxt,k)
        #os.system(cmd)




